package com.winit.schedule.job.tracker.repository.mongo;

import java.util.ArrayList;
import java.util.List;

import com.google.code.morphia.query.Query;
import com.google.code.morphia.query.UpdateOperations;
import com.winit.schedule.core.model.JobID;
import com.winit.schedule.core.model.TaskID;
import com.winit.schedule.core.repository.TaskRepository;
import com.winit.schedule.core.repository.constants.TaskStatus;
import com.winit.schedule.core.repository.entity.JobEntity;
import com.winit.schedule.core.repository.entity.TaskEntity;
import com.winit.schedule.core.repository.query.TaskQuery;
import com.winit.schedule.store.mongo.AbstractMongoRepository;

public class MongoTaskRepository extends AbstractMongoRepository<TaskEntity> implements TaskRepository {

    @Override
    public TaskEntity saveTaskEntity(TaskEntity entity) {
        Query<TaskEntity> query = createQuery().field("taskId.taskId").equal(entity.getTaskId().getTaskId());
        UpdateOperations<TaskEntity> ops = ds.createUpdateOperations(TaskEntity.class)
            .set("taskGroup", entity.getTaskGroup())
            .set("taskName", entity.getTaskName())
            .set("autoRecovery", entity.isAutoRecovery())
            .set("immediatelyRun", entity.isImmediatelyRun())
            .set("maxThreads", entity.getMaxThreads())
            .set("sleepTimeNoData", entity.getSleepTimeNoData())
            .set("eachFetchDataNum", entity.getEachFetchDataNum())
            .set("sleepTimeInterval", entity.getSleepTimeInterval())
            .set("threadsOfRegion", (entity.getThreadsOfRegion() <= 0 ? 1 : entity.getThreadsOfRegion()));
        if (null != entity.getInterfaceType()) {
            ops.set("interfaceType", entity.getInterfaceType());
        }

        if (null != entity.getRunStartTime()) {
            ops.set("runStartTime", entity.getRunStartTime());
        }
        if (null != entity.getRunEndTime()) {
            ops.set("runEndTime", entity.getRunEndTime());
        }
        if (null != entity.getTaskItem()) {
            ops.set("taskItem", entity.getTaskItem());
        }
        if (null != entity.getIpWhiteList()) {
            ops.set("ipWhiteList", entity.getIpWhiteList());
        }
        if (null != entity.getTaskParameter()) {
            ops.set("taskParameter", entity.getTaskParameter());
        }
        if (null != entity.getTaskStatus()) {
            ops.set("taskStatus", entity.getTaskStatus());
        }
        if (null != entity.getTaskItem()) {
            ops.set("taskItem", entity.getTaskItem());
        }

        return ds.findAndModify(query, ops, false, true);
    }

    @Override
    public List<TaskEntity> saveTaskEntity(List<TaskEntity> taskEntitys) {

        if (null != taskEntitys && taskEntitys.size() > 0) {
            for (TaskEntity entity : taskEntitys) {
                this.saveTaskEntity(entity);
            }
        }
        return taskEntitys;
    }

    @Override
    public void updateTaskInterfaceType(List<TaskEntity> taskEntitys) {
        if (null != taskEntitys && taskEntitys.size() > 0) {
            for (TaskEntity entity : taskEntitys) {
                this.updateTaskDefineInfo(entity);
            }
        }
    }

    private void updateTaskDefineInfo(TaskEntity entity) {
        Query<TaskEntity> query = createQuery().field("taskId.taskId").equal(entity.getTaskId().getTaskId());
        TaskEntity task = query.get();
        UpdateOperations<TaskEntity> ops = ds.createUpdateOperations(TaskEntity.class)
            .set("taskGroup", entity.getTaskGroup())
            .set("taskName", entity.getTaskName())
            .set("interfaceType", entity.getInterfaceType());

        if (null == task) {
            ops.set("taskId", entity.getTaskId());
        }
        ds.findAndModify(query, ops, false, true);
    }

    @Override
    public TaskEntity getTaskEntityById(String taskId) {
        return createQuery().field("taskId.taskId").equal(taskId).get();
    }

    @Override
    public void removeTaskEntity(TaskID taskId) {
        Query<TaskEntity> query = createQuery().field("taskId.taskId").equal(taskId.getTaskId());
        ds.delete(query);
    }

    @Override
    public List<TaskEntity> findTaskEntity(TaskQuery entity) {
        Query<TaskEntity> query = createQuery();

        if (null != entity) {
            query.offset(entity.getOffset());
            if (entity.getLimit() > 0) {
                query.limit(entity.getLimit());
            }
        }

        return query.asList();
    }

    @Override
    public TaskEntity updateTaskStatus(String taskId, TaskStatus status) {
        Query<TaskEntity> query = createQuery().field("taskId.taskId").equal(taskId);
        UpdateOperations<TaskEntity> ops = ds.createUpdateOperations(TaskEntity.class).set("taskStatus", status);
        return ds.findAndModify(query, ops);
    }

    @Override
    public List<TaskID> queryAllStopTaskIDs() {
        Query<TaskEntity> query = createQuery().field("taskStatus").equal(TaskStatus.STOP);
        List<TaskEntity> tasks = query.asList();
        List<TaskID> ids = new ArrayList<TaskID>();
        for (TaskEntity task : tasks) {
            ids.add(task.getTaskId());
        }
        return ids;
    }

    @Override
    public List<String> queryAllTaskIDByJobId(JobID jobId) {
        Query<TaskEntity> query = createQuery().field("taskId.jobId").equal(jobId);
        List<TaskEntity> tasks = query.asList();
        List<String> ids = new ArrayList<String>();
        for (TaskEntity task : tasks) {
            ids.add(task.getTaskId().getTaskId());
        }
        return ids;
    }

    @Override
    public void updateTaskStatus(JobID jobId, TaskStatus status) {
        Query<TaskEntity> query = createQuery().field("taskId.jobId").equal(jobId);
        UpdateOperations<TaskEntity> ops = ds.createUpdateOperations(TaskEntity.class).set("taskStatus", status);
        ds.update(query, ops);
    }

    @Override
    public List<TaskEntity> queryTasksNotRelatedJob() {
        Query<TaskEntity> query = createQuery().field("taskId.jobId").doesNotExist().field("taskStatus").exists();
        return query.asList();
    }

    @Override
    public void saveJobIdForTask(JobEntity job, List<String> taskIds) {
        Query<TaskEntity> query = createQuery().field("taskId.taskId").in(taskIds);
        UpdateOperations<TaskEntity> ops = ds.createUpdateOperations(TaskEntity.class);
        ops.set("taskId.jobId", job.getJobId());
        ops.set("runStartTime", job.getRunStartTime());
        ops.set("runEndTime", null != job.getRunEndTime() ? job.getRunEndTime() : "");
        ops.set("immediatelyRun", job.isImmediatelyRun());
        ds.update(query, ops);
    }

    @Override
    public void delJobIdForTask(JobID jobId) {
        Query<TaskEntity> query = createQuery().field("taskId.jobId").equal(jobId);
        UpdateOperations<TaskEntity> ops = ds.createUpdateOperations(TaskEntity.class);
        ops.unset("taskId.jobId");
        ops.unset("runStartTime");
        ops.unset("runEndTime");
        ds.update(query, ops);

    }

    @Override
    public List<TaskEntity> queryTasksRelatedJob(JobID jobId) {
        Query<TaskEntity> query = createQuery().field("taskId.jobId").equal(jobId);
        return query.asList();
    }

    @Override
    public long queryUnfinishedTaskNumber(String jobId) {
        Query<TaskEntity> query = createQuery().field("taskId.jobId")
            .equal(jobId)
            .field("taskStatus")
            .notEqual(TaskStatus.COMPLETE);
        return query.countAll();
    }
}
